import os
from os.path import join, isfile
import sys
sys.path.append("..")
import numpy as np
import torch
from torch import nn
import pytorch3d.io as IO
import trimesh
from utils.VTS_object import get_obj_info
from smplx.smplx.lbs import batch_rodrigues
from transforms3d.axangles import mat2axangle
import open3d as o3d


class Object_retargeting(nn.Module):
    def __init__(self, init_params, cfg):
        super(Object_retargeting, self).__init__()
        self.rotation = nn.Parameter(init_params['rotation'].clone().detach(), requires_grad=True)
        self.translation = nn.Parameter(init_params['translation'].clone().detach(), requires_grad=True)
        self.cfg = cfg
    def forward(self):
        result = {
            'rotation': self.rotation,
            'translation': self.translation,
        }
        return result


def obj_retargeting_new_obj(data_path, source_mesh_path, target_mesh_path, start_frame, end_frame, device='cuda:0', q_mat=None):
    object_dir = join(data_path, 'aligned_objposes.npy')
    # 301 x 4 x 4
    source_object_result = np.load(object_dir, allow_pickle=True)
    # object_mesh = trimesh.load_mesh(mesh_path)
    source_object_mesh = IO.load_objs_as_meshes([source_mesh_path])
    source_vert = source_object_mesh.verts_packed().to(torch.float64).to(device)  # shape = (N, 3)

    target_object_mesh = IO.load_objs_as_meshes([target_mesh_path])
    target_vert = target_object_mesh.verts_packed().to(torch.float64).to(device)  # shape = (N, 3)
    new_obj_mesh = trimesh.Trimesh(vertices=target_vert.detach().cpu().numpy(), faces=target_object_mesh.faces_packed().detach().cpu().numpy())

    if q_mat is not None:
        q_mat = q_mat.type(torch.float64).to(device)
        new_mesh_vert = torch.matmul(target_vert, q_mat[:3, :3].transpose(1, 0)) + q_mat[:3, 3]
        new_obj_mesh = trimesh.Trimesh(vertices=new_mesh_vert.detach().cpu().numpy(), faces=target_object_mesh.faces_packed().detach().cpu().numpy())

    if (len(source_object_result) < end_frame):
        end_frame = len(source_object_result)
        print("Warning: end_frame is larger than the length of object_result, set end_frame to {}".format(str(end_frame)))

    rot_vec = []
    for i in range(source_object_result.shape[0]):
        d, a = mat2axangle(source_object_result[i][:3, :3])
        rot_vec.append(d * a)
    trans_vec = source_object_result[:, :3, 3]

    init_params = {
        "rotation": torch.tensor(rot_vec, dtype=torch.float64).detach().to(device),
        "translation": torch.from_numpy(trans_vec).detach().to(device),
    }

    M = end_frame - start_frame
    N_epoch = 1200
    result_params = []
    for i in range(start_frame, end_frame, M):
        print("------------ optimizing frame {} - {} -----------------------".format(str(i), str(i+M-1)))
        ids = np.arange(i, min(i+M, end_frame))
        L = len(ids)
        sequence_init_params = {}
        # sequence_init_params["obj_pose"] = init_params["obj_pose"][ids].clone()
        for key in init_params.keys():
            sequence_init_params[key] = init_params[key][ids].clone()
        cfg = {}
        optim_model = Object_retargeting(sequence_init_params, cfg)
        optim_model.to(device)
        optimizer = torch.optim.Adam(optim_model.parameters(), lr=0.005)
        optim_model.train()
        for epoch in range(N_epoch):
            optimizer.zero_grad()
            result = optim_model()
            # print(result)
            re_rot_vec = result['rotation'].type(torch.float64)
            re_trans_vec = result["translation"].type(torch.float64)
            
            rot_loss = torch.mean(torch.abs(re_rot_vec - init_params["rotation"][ids]))

            trans_loss = torch.mean(torch.abs(re_trans_vec - init_params["translation"][ids]))
            re_rot_mat = batch_rodrigues(re_rot_vec)
            ori_rot_mat = batch_rodrigues(init_params["rotation"][ids])
            origin_vert = torch.matmul(ori_rot_mat[:, ], source_vert.transpose(1, 0)).permute(0, 2, 1) + init_params["translation"][ids][:, ].unsqueeze(1)
            mesh_vert = torch.matmul(re_rot_mat[:, ], target_vert.transpose(1, 0)).permute(0, 2, 1) + re_trans_vec[:, ].unsqueeze(1)
            bottom_loss = torch.sum((origin_vert[:, :, 1].min() - mesh_vert[:, :, 1].min()) ** 2)
            
            y_loss = torch.sum(mesh_vert[:, :, 1].clamp(None, 0)**2)  # liuyun: 希望让这个loss影响更多的vertex
            # print(mesh_vert.shape)
            # smoothness_rot = torch.sum((2 * re_rot_vec[1:-1,] - re_rot_vec[2:,] - re_rot_vec[:-2, ]) ** 2)
            smoothness_rot = torch.sum((2 * re_rot_mat[1:-1] - re_rot_mat[2:] - re_rot_mat[:-2]) ** 2)
            smoothness_trans = torch.sum((2 * re_trans_vec[1:-1] - re_trans_vec[2:] - re_trans_vec[:-2]) ** 2)
            smoothness = smoothness_rot + smoothness_trans
            
            lambda_rot_loss, lambda_trans_loss, lambda_y_loss, lambda_smoothness, lambda_bottom_loss = 500, 0.005, 0.01, 1, 100
            loss = lambda_rot_loss * rot_loss + lambda_trans_loss * trans_loss + lambda_y_loss * y_loss + lambda_smoothness * smoothness + lambda_bottom_loss * bottom_loss
            # loss = 0.1 * pnt_loss
            
            if epoch % 100 == 0 or (epoch == N_epoch - 1):
                print(epoch, loss.item(), lambda_rot_loss * rot_loss.item(), lambda_trans_loss * trans_loss.item(), lambda_y_loss * y_loss.item(), lambda_smoothness * smoothness.item(), lambda_bottom_loss * bottom_loss.item())
            loss.backward()
            optimizer.step()
        
        optim_model.eval()
        result = optim_model()
        for j in range(L):
            result_params.append({"rotation":result["rotation"][j:j+1].detach(),
                                  "translation":result["translation"][j:j+1].detach()})
        
        init_params["rotation"][ids] = result["rotation"].clone().detach()
        init_params["translation"][ids] = result["translation"].clone().detach()
    
    result = {
        "obj_mesh": new_obj_mesh,
        "obj_poses": result_params,
    }
    return result



def obj_retargeting(data_path, mesh_path, q_mat, start_frame, end_frame, device='cuda:0', new_obj=False):
    object_dir = join(data_path, 'aligned_objposes.npy')
    # 301 x 4 x 4
    object_result = np.load(object_dir, allow_pickle=True)
    # object_mesh = trimesh.load_mesh(mesh_path)
    object_mesh = IO.load_objs_as_meshes([mesh_path])
    vert = object_mesh.verts_packed().to(torch.float64)  # shape = (N, 3)
    # q_mat = q_mat.to(device).type(torch.float64)
    q_mat = q_mat.type(torch.float64).to(device)
    new_mesh_vert = torch.matmul(vert.to(device), q_mat[:3, :3].transpose(1, 0)) + q_mat[:3, 3]
    new_obj_mesh = trimesh.Trimesh(vertices=new_mesh_vert.detach().cpu().numpy(), faces=object_mesh.faces_packed().detach().cpu().numpy())
    
    if (len(object_result) < end_frame):
        end_frame = len(object_result)
        print("Warning: end_frame is larger than the length of object_result, set end_frame to {}".format(str(end_frame)))

    # rot_vec = [mat2axangle(object_result[i][:3, :3])[0] * mat2axangle for i in range(object_result.shape[0])]
    rot_vec = []
    for i in range(object_result.shape[0]):
        d, a = mat2axangle(object_result[i][:3, :3])
        rot_vec.append(d * a)
    trans_vec = object_result[:, :3, 3]

    
    init_params = {
        "rotation": torch.tensor(rot_vec, dtype=torch.float64).detach().to(device),
        "translation": torch.from_numpy(trans_vec).detach().to(device),
    }
    
    # liuyun:
    # 每一帧object pose的解空间是 3D的axis-angle(也叫rotation vector) + 3D的translation vector
    # 用batch_rodrigues(在第22行import了)把(N,3)的batched axis-angle可微地转换成(N, 3, 3)的batched rotation matrix
    # loss 可以是rotation matrix的L1或L2 loss, 也可以是axis-angle的L1/L2 loss
    # 把 gt rotation matrix转化成gt axis-angle (numpy操作, 不需支持可微): mat2axangle(在第23行import了)
    
    M = 50
    N_epoch = 1000
    result_params = []
    for i in range(start_frame, end_frame, M):
        print("------------ optimizing frame {} - {} -----------------------".format(str(i), str(i+M-1)))
        ids = np.arange(i, min(i+M, end_frame))
        L = len(ids)
        sequence_init_params = {}
        # sequence_init_params["obj_pose"] = init_params["obj_pose"][ids].clone()
        for key in init_params.keys():
            sequence_init_params[key] = init_params[key][ids].clone()
        cfg = {}
        optim_model = Object_retargeting(sequence_init_params, cfg)
        optim_model.to(device)
        optimizer = torch.optim.Adam(optim_model.parameters(), lr=0.005)
        optim_model.train()
        for epoch in range(N_epoch):
            optimizer.zero_grad()
            result = optim_model()
            # print(result)
            re_rot_vec = result['rotation'].type(torch.float64)
            re_trans_vec = result["translation"].type(torch.float64)
            
            rot_loss = torch.mean(torch.abs(re_rot_vec - init_params["rotation"][ids]))

            trans_loss = torch.mean(torch.abs(re_trans_vec - init_params["translation"][ids]))
            # pnt_loss = max(0, -torch.mean(torch.min(curr_bb[:, :, 1])))
            re_rot_mat = batch_rodrigues(re_rot_vec)
            mesh_vert = torch.matmul(re_rot_mat[:, ], new_mesh_vert.transpose(1, 0)).permute(0, 2, 1) + re_trans_vec[:, ].unsqueeze(1)
            # y_loss = torch.max(-torch.min(mesh_vert[:, :, 1]), torch.zeros_like(torch.min(mesh_vert[:, :, 1])))
            y_loss = torch.sum(mesh_vert[:, :, 1].clamp(None, 0)**2)  # liuyun: 希望让这个loss影响更多的vertex
            # print(mesh_vert.shape)
            # smoothness_rot = torch.sum((2 * re_rot_vec[1:-1,] - re_rot_vec[2:,] - re_rot_vec[:-2, ]) ** 2)
            smoothness_rot = torch.sum((2 * re_rot_mat[1:-1] - re_rot_mat[2:] - re_rot_mat[:-2]) ** 2)
            smoothness_trans = torch.sum((2 * re_trans_vec[1:-1] - re_trans_vec[2:] - re_trans_vec[:-2]) ** 2)
            smoothness = smoothness_rot + smoothness_trans
            
            lambda_rot_loss, lambda_trans_loss, lambda_y_loss, lambda_smoothness = 500, 0.005, 0.01, 1
            loss = lambda_rot_loss * rot_loss + lambda_trans_loss * trans_loss + lambda_y_loss * y_loss + lambda_smoothness * smoothness
            # loss = 0.1 * pnt_loss
            
            if epoch % 100 == 0 or (epoch == N_epoch - 1):
                print(epoch, loss.item(), lambda_rot_loss * rot_loss.item(), lambda_trans_loss * trans_loss.item(), lambda_y_loss * y_loss.item(), lambda_smoothness * smoothness.item())
            loss.backward()
            optimizer.step()
        
        optim_model.eval()
        result = optim_model()
        for j in range(L):
            result_params.append({"rotation":result["rotation"][j:j+1].detach(),
                                  "translation":result["translation"][j:j+1].detach()})
        
        init_params["rotation"][ids] = result["rotation"].clone().detach()
        init_params["translation"][ids] = result["translation"].clone().detach()
    
    result = {
        "obj_mesh": new_obj_mesh,
        "obj_poses": result_params,
    }
    return result


if __name__ == '__main__':
    
    #########################################################################################
    obj_dataset_dir = "/share/datasets/HHO_object_dataset_final"
    # data_dir = '/share/datasets/HHO_dataset/data/20230701/012/'
    data_dir = '/share/datasets/HHO_dataset/data/20230806_1/012/'
    q_mat = torch.from_numpy(np.array([[1.2, 0, 0, 0], [0, 2, 0, 0], [0, 0, 1.5, 0], [0, 0, 0, 1]]))
    start_frame, end_frame = 0, 300
    device = "cuda:0"
    #########################################################################################
    
    obj_name, obj_data_path = get_obj_info(data_dir, obj_dataset_dir)
    print(obj_name, obj_data_path)
    assert isfile(obj_data_path)
    
    obj_retargeting_result = obj_retargeting(data_dir, obj_data_path, q_mat, start_frame=start_frame, end_frame=end_frame, device=device)
    np.save('0806_1_012_12215.npy', obj_retargeting_result["obj_poses"])
    # print(result_params['obj_pose'])
    